Wheat Head Detection using Pyspark, Python, VGGUnet

Author : Deepika Sharma Date : June 2020 - July 2020

Folder structure:

  1. Input(train) data 3422 images.
  2. Train csv : CSV containing information of bounding boxes around wheat spikes in each image.
  3. Test : Images on which prediction has to be made, prediction output has to be in the format present in submission csv.
  4. Submission csv : Format in which submission has to be made. It has 2 columns : 1. image name, 2. probability of bounding box containing wheat spike, correspnding co-orinates x,y,h,w ... prob and x,y,h,w for each predicted bbox.
  5. Intermediate folder (temporary folder to keep images to be processed batchwise). Batchsize needs to be changed based on your system specifications.

Steps:

Step 1: In this kernal, I'll be using pyspark for data preparation as it requires manipulating 3000 images (train size). ~300 for kept for testing. Train csv does not contain information for all images present in Raw data, pls do not worry if numbers do not sum up.

For training, each image is converted in to binary image, where 0's represent background, 1's represent wheat spikes Encoding is done using bounding boxes coordinates present in the Train csv.

Step 2: Model (Keras vgg_unet) training on ~3000 images with n_classes=2 , input_height=1024, input_width=1024.

images are resized to reduced computation on CPU.

Step 3: Prediction on ~400 images to remove background (0) and only predict wheat spikes as foreground (1)
Step 4: Blob detection using blob_dog (difference of Gaussian) to detect blobs in model predicted segments for wheat spikes.
Step 5: IoU calculation, and computation of probability of each detection in each image.

Step 1: Data Preparation

Import libraries

In [1]:
import cv2
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


### Scikit learn libraries ### 
from skimage.color import rgb2gray
from skimage.feature import blob_dog
from keras.models import model_from_json


### Pyspark libraries ###
import os,gc
import shutil
import pyspark
import itertools
from operator import add
from pyspark.sql.types import *
from pyspark import SparkContext
from pyspark.sql import DataFrame
from pyspark.sql.functions import udf
from pyspark.ml.image import ImageSchema
Using TensorFlow backend.

Reading input files

In [2]:
file = pd.read_csv("train_csv/train.csv")  # This file contains bbox co-ordinates of wheat spikes for each image.

# Train csv has multiple rows for each image depending upon number of wheat spikes in each image.
ids = [file.iloc[i][0] for i in range(file.shape[0])] 

# Factor by which we will resize each image as each raw image has size 1024*1024, which will be too large for network to train on to.
factor = 2; size = 512; batch_size = 500; x_axis_gaps = 5; coords_dist_th = 15; N =2999 ;N_ = N+30

Annotated samples

In [3]:
sample = os.listdir("train")[0:3]

for img_name in sample:
    image = cv2.imread('train/'+ img_name)
    
    # Get image ids for each image to access the segmentation co-ordinates
    image_ids = [l for l,val in enumerate(ids) if val == str(img_name.split(".")[0])]
    list_of_coords = [[int(float(val)) for val in file.iloc[v][3][1:len(file.iloc[v][3])-1].split(",")] for v in image_ids] 
            
    for l in list_of_coords:        
        image_annotated = cv2.rectangle(image,(int(l[0]),int(l[1])),(int((l[0]+l[2])),int((l[1]+l[3]))),(255,0,0),4)

    fig, axes = plt.subplots(1, 2, figsize=(10, 5))#, sharex=True, sharey=True)
    ax = axes.ravel()
    image = cv2.imread('train/'+ img_name)    
    ax[0].imshow(cv2.cvtColor(image,cv2.COLOR_RGB2BGR))
    ax[1].imshow(cv2.cvtColor(image_annotated,cv2.COLOR_RGB2BGR))
        
    plt.show()

Let's get started with Data Preparation ................

Converting images to binary for training

  • Data Preparation is the most critical part. For which, I had written a seperate module and imported the same here.
  • Usually, our notebooks get consumed by "many Data scientists (referenced ones :)" , hence leaving data preparation part away to keep personal intellectual property safe.
In [4]:
if not os.path.exists('train_annotated'):
        os.makedirs('train_annotated')
        
from data_prep_module import data_prep

Let's Spark it up!

Setting up Spark Context

Note, I made a few chnages in the backend code for Pyspark UDF to work here. In case you face challanges, feel free to comment/connect.

In [5]:
def run_spark():   
    sc=SparkContext(master="local[15]") # number[15] can be changed based on your system specification.
#     print(sc.binaryFiles("spark_temp/*.png"))
    image_df = sc.binaryFiles("spark_temp") # Reading images into Spark context
    
    pyspark.sql.udf.UDFRegistration.register(name="data_prep", f = data_prep, returnType=StringType()) #registering UDF 
    # such that Spark context recongises the function used for data preparation.
    
    job = [data_prep(ids,factor,file,size,x[0].split("/")[-1]) for idx,x in enumerate(image_df.take(batch_size))]
    shutil.rmtree('spark_temp') #remove the folder once 'one batch' is complete to avoid Spark remembering 
    # indices for images it is done with.
    sc.stop() # stop the spark context else Spark will have unnecessary information cached whcih we do not require anymore.
    gc.collect() # remove any other cache which memory might have been holding.
    
    
In [6]:
input_list = os.listdir("train")[0:N]
input_seq = sorted(list(set([val for val in range(0, len(input_list),batch_size)]+ [len(input_list)])))

def image_transfer(value):
    if not os.path.exists('spark_temp'):
        os.makedirs('spark_temp')
    image = cv2.imread('train/'+value)
    cv2.imwrite("spark_temp/" + value.split(".")[0]+".jpg",cv2.resize(image,(size,size)))
    
    
for val in range(len(input_seq)-1):
    img_batch = input_list[input_seq[val]:input_seq[val]+batch_size]
    [image_transfer(val) for val in img_batch]
  
    run_spark()  
    

Data Preparation is Done!

Step 2: Model Training

Lets' do a bit checking on our data.

In [7]:
if not os.path.exists('train_temp'):
        os.makedirs('train_temp')
def image_transfer_(value):    
    
    cv2.imwrite("train_temp/" + value.split(".")[0]+".png",cv2.resize(cv2.imread('train/'+value.split(".")[0]+".jpg"),(size,size)))
    
first_3000 = os.listdir("train_annotated")
task = [image_transfer_(val) for val in first_3000]
In [8]:
error_list = [val if np.max(cv2.imread('train_annotated/'+val)) > 1 else None for val in os.listdir("train_annotated")]
error_list = [val for val in error_list if val != None]
if len(error_list) > 0:
    print(error_list)
    [os.remove("train_temp/"+val.split(".")[0]+".jpg") for val in error_list]
    [os.remove("train_annotated/"+val) for val in error_list]
In [9]:
# import tensorflow.compat.v1 as tensorflow
from keras_segmentation.models.unet import vgg_unet
model = vgg_unet(n_classes=2, input_height=size, input_width=size)
model.train(
    train_images =  "train_temp/",#train
    train_annotations = "train_annotated/",n_classes = 2,epochs=10,steps_per_epoch=5,#annotations_prepped_train_v3
)
WARNING:tensorflow:From D:\Anaconda\envs\wheat_head_count\lib\site-packages\tensorflow\python\ops\resource_variable_ops.py:1666: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
Verifying training dataset
100%|██████████████████████████████████████████████████████████████████████████████| 2999/2999 [04:36<00:00, 10.83it/s]
Dataset verified! 
Epoch 1/10
5/5 [==============================] - 130s 26s/step - loss: 0.9549 - acc: 0.5696
Epoch 2/10
5/5 [==============================] - 118s 24s/step - loss: 0.8326 - acc: 0.5952
Epoch 3/10
5/5 [==============================] - 123s 25s/step - loss: 0.6160 - acc: 0.6943
Epoch 4/10
5/5 [==============================] - 123s 25s/step - loss: 0.5958 - acc: 0.7141
Epoch 5/10
5/5 [==============================] - 134s 27s/step - loss: 0.5734 - acc: 0.7277
Epoch 6/10
5/5 [==============================] - 121s 24s/step - loss: 0.6013 - acc: 0.7269
Epoch 7/10
5/5 [==============================] - 117s 23s/step - loss: 0.5192 - acc: 0.7638
Epoch 8/10
5/5 [==============================] - 111s 22s/step - loss: 0.5785 - acc: 0.7141
Epoch 9/10
5/5 [==============================] - 114s 23s/step - loss: 0.4876 - acc: 0.7875
Epoch 10/10
5/5 [==============================] - 111s 22s/step - loss: 0.4932 - acc: 0.7829
In [10]:
 #saving model to disk
from keras.models import model_from_json
model_json = model.to_json()

with open("model.json", "w") as json_file:
    json_file.write(model_json)
model.save_weights("model.h5")

Step 3: Segmentation prediction

Step 4: Blob Detection

Step 5: IoU calculation

In [11]:
for val in os.listdir("test/"):    
    cv2.imwrite("test/"+val.split(".")[0]+".jpg",cv2.resize(cv2.imread('test/'+val.split(".")[0]+".jpg"),(size,size)))
In [12]:
import warnings
warnings.simplefilter("ignore")
accuracy = []; submission = {"image_id" : [], "width" : [], "height" : [],"bbox" : [], "Prob":[]}

for idx,img_name in enumerate(os.listdir("test/")[0:5]):
    if not os.path.exists('segmentation'):
        os.makedirs('segmentation')
    if not os.path.exists('Output'):
        os.makedirs('Output')
        
    
    out = model.predict_segmentation(
    inp="test/"+ img_name,
    out_fname="segmentation/" + img_name.split(".")[0] + ".png" 
    )  
    
    out_image = cv2.imread('segmentation/'+ img_name.split(".")[0] + ".png")
    _image = cv2.imread('test/'+ img_name.split(".")[0] + ".jpg")
    
    ## Step 4:Blob detection
    image_gray = rgb2gray(out_image)
    
    blobs_dog = blob_dog(image_gray, max_sigma=30, threshold=.10)

    try:
        blobs_dog[:, 2] = blobs_dog[:, 2] * math.sqrt(2)
        submission["image_id"].append(img_name.split(".")[0])
        submission["width"].append(size); submission["height"].append(size)
        submission["bbox"].append([str([int(blob[0]), int(blob[1]),int(blob[2]), int(blob[2])]) for blob in blobs_dog])
        submission["Prob"].append(.50)
    except:
        submission["image_id"].append(img_name.split(".")[0])
        submission["width"].append(size); submission["height"].append(size)
        submission["bbox"].append([])
        submission["Prob"].append(1)

    
    blob_list_sorted = sorted([[int(blob[0]), int(blob[1]),int(blob[2])] for blob in blobs_dog])
    len_val = len(blob_list_sorted)
    for i,val in enumerate(range(0,len(blob_list_sorted))):
        for j in range(i+1,len(blob_list_sorted)-2):
            if max(i,j) < len_val and max([abs(np.diff(x)) for x in zip(blob_list_sorted[j][0:2], blob_list_sorted[i][0:2])]) < coords_dist_th:
                blob_list_sorted = blob_list_sorted+[[min(x) for x in zip(blob_list_sorted[j], blob_list_sorted[i])][0:2]+[blob_list_sorted[j][2]+blob_list_sorted[i][2]]]
                del blob_list_sorted[i]
                del blob_list_sorted[j]  
                len_val = len(blob_list_sorted)
    try:
        for blob in blobs_dog:                   
            y, x, r = blob
            Iou = cv2.rectangle(_image,(int(x),int(y)),(int((x+r)),int((y+r))),(255,0,0),2)

#         print(Iou)
        cv2.imwrite("Output/" + img_name.split(".")[0] + ".png",Iou)
        fig, axes = plt.subplots(1, 2, figsize=(12, 6))#, sharex=True, sharey=True)
        ax = axes.ravel()

        ax[0].imshow(cv2.cvtColor(cv2.imread('test/'+ img_name.split(".")[0] + ".jpg"),cv2.COLOR_RGB2BGR))
        ax[1].imshow(Iou)
   
        plt.show()
    except:pass
pd.DataFrame(submission).to_csv("submission.csv")    
In [13]:
pd.DataFrame(submission).to_csv("submission.csv") 

Results are looking alryt! Let's test on few more images..

In [14]:
if not os.path.exists('train_50'):
        os.makedirs('train_50')
for val in os.listdir("train/")[N:N_]:    
    cv2.imwrite("train_50/"+val.split(".")[0]+".png",cv2.resize(cv2.imread('train/'+val.split(".")[0]+".jpg"),(size,size)))
In [15]:
warnings.simplefilter("ignore")

for idx,img_name in enumerate(os.listdir("train_50")):
    if not os.path.exists('segmentation'):
        os.makedirs('segmentation')
    
    
    out = model.predict_segmentation(
    inp="train_50/"+ img_name,
    out_fname="segmentation/" + img_name.split(".")[0] + ".png" 
    )  
    
    out_image = cv2.imread('segmentation/'+ img_name.split(".")[0] + ".png")
    _image = cv2.imread('train_50/'+ img_name.split(".")[0] + ".png")


    ## Step 4:Blob detection
    image_gray = rgb2gray(out_image)
    blobs_dog = blob_dog(image_gray, max_sigma=30, threshold=.10)
    
    blob_list_sorted = sorted([[int(blob[0]), int(blob[1]),int(blob[2])] for blob in blobs_dog])
    len_val = len(blob_list_sorted)
    for i,val in enumerate(range(0,len(blob_list_sorted))):
        for j in range(i+1,len(blob_list_sorted)-2):
            if max(i,j) < len_val and max([abs(np.diff(x)) for x in zip(blob_list_sorted[j][0:2], blob_list_sorted[i][0:2])]) < coords_dist_th:
                blob_list_sorted = blob_list_sorted+[[min(x) for x in zip(blob_list_sorted[j], blob_list_sorted[i])][0:2]+[blob_list_sorted[j][2]+blob_list_sorted[i][2]]]
                del blob_list_sorted[i]
                del blob_list_sorted[j]  
                len_val = len(blob_list_sorted)
   
    try:
        for blob in blobs_dog:                   
            y, x, r = blob
            Iou = cv2.rectangle(_image,(int(x),int(y)),(int((x+r)),int((y+r))),(255,0,0),2)         
        
        fig, axes = plt.subplots(1, 2, figsize=(12, 6))
        ax = axes.ravel()

        ax[0].imshow(cv2.cvtColor(cv2.imread('train_50/'+ img_name.split(".")[0] + ".png"),cv2.COLOR_RGB2BGR))
        ax[1].imshow(Iou)

        plt.show()
        
    except:pass

End Notes

  • There are classes in image dataset, few images with too much sunlight, few images with clear distinction between wheat heads and background.
  • It will be required for me to train the model with additional sophistication to be able to predict wheat heads in different light conditions, different orientation due to wind and color variation depending on the age of the crop.
  • Look forward to next notebook on the same..

To all my fellow Data Scientists, non privileged ones, privileged ones, in the industry because of relatives, or by yourself... Happy coding ;)

This work took me a month time to identify right pre-processing technique which worked for all (almost) images. if you notice, results are pretty satisfactory in terms of algorithm picking all the wheat heads in an image despite the varied size of the wheat heads, different light conditions and orientation.

I agree, there is scope for improvement (which I will take up later in sometime..) but I am equally proud that the solution is able to seperate wheat heads from leaves which are of same color and at times, same shape and size. and top that solution is able to discard the background which is soil, dry leaves and etc..

Few Applications of this exercise:

Appreciate the effort.., do critic, share feedback. You all are most welcome.

In [ ]: